#!/usr/bin/env python

import functools
import math
import sys

from matplotlib import pyplot


def fxrange(start, end, step):
    x = start
    while x <= end:
        yield x
        x += step


def choose_step(method, eps):
    h = 1.0
    x_vals_1, y_vals_1 = method(h)
    while h:
        h /= 2
        x_vals_2, y_vals_2 = method(h)
        shorten_y_vals = [y for i, y in enumerate(y_vals_2) if i % 2 == 0]
        dy_vals = [math.fabs(y1 - y2) for y1, y2 in zip(y_vals_1, shorten_y_vals)]
        dy_vals = [True if dy < eps else False for dy in dy_vals]
        if all(dy_vals):
            return h
        x_vals_1, y_vals_1 = x_vals_2, y_vals_2
    # if we could not find step good enough
    return sys.float_info.min

def solve_tridiagonal_system(A):
    n = len(A)
    a_coeffs = [0]*(n+1)
    b_coeffs = [0]*(n+1)
    for i in xrange(1, n+1):
        a_coeffs[i] = - A[i-1][2] / (A[i-1][0]*a_coeffs[i-1] + A[i-1][1])
        b_coeffs[i] = (A[i-1][3] - A[i-1][0]*b_coeffs[i-1]) / (A[i-1][0]*a_coeffs[i-1] + A[i-1][1])
    x_values = [0]*n
    x_values[n-1] = b_coeffs[n]
    for i in xrange(n-1, 0, -1):
        x_values[i-1] = a_coeffs[i] * x_values[i] + b_coeffs[i]
    return x_values


def aprox(x0, xn, a, b, h):
    """
    y''(xk) ~= (yk+1 - 2*yk + yk-1) / h**2      k = [1,n-1]
    y(a) = A
    y(b) = B
    a * (y''k) / h**2  +  (1+b*xk**2)*yk = -1
    a * y''k + (1+b*xk**2) * yk * h**2 = -h**2
    a * yk-1 - ((1+b*xk**2) * h**2 - 2*a) * yk + a*yk+1 = -h**2
    coefficients:
    a           b                               c           d
    a           ((1 + b*xk**2)*h**2 - 2*a)      a           -h**2
    """

    x_values = list(fxrange(x0, xn, h))
    n = len(x_values) - 1
    A = []

    for i in xrange(1, n):
        A.append([a, ((1 + b*x_values[i]**2) * h**2 - 2*a), a, -h**2])
    A[0][0] = 0
    A[-1][-2] = 0
    y_values = solve_tridiagonal_system(A)
    # assuming that y(a) = y(b) = 0
    y_values = [0] + y_values + [0]

    return x_values, y_values


def main():
    a = math.sin(17)
    b = math.sin(17)
    eps = 0.001

    h = choose_step(functools.partial(aprox, -1, 1, a, b), eps)
    print 'step:', h

    x_values, y_values = aprox(-1, 1, a, b, h)

    for i in xrange(len(x_values)):
        print x_values[i], y_values[i]

    pyplot.plot(x_values, y_values)
    pyplot.show()


if __name__ == '__main__':
    main()
